# train_ensemble.py
import numpy as np
import logging
import pickle
import time
import os
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, VotingClassifier, BaggingClassifier, AdaBoostClassifier, GradientBoostingClassifier
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score
from util import createXY
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def train_ensemble_models():
    """训练多种集成学习模型并选择最佳模型"""
    
    # 使用正确的数据路径
    train_folder = "./data/train"
    
    logging.info(f"使用数据目录: {train_folder}")
    logging.info("读取所有图像，生成X和y")
    
    try:
        X, y = createXY(train_folder=train_folder, dest_folder=".", method='flat')
        X = np.array(X).astype('float32')
        y = np.array(y)
        
        logging.info(f"X.shape: {X.shape}")
        logging.info(f"y.shape: {y.shape}")
        
        if X.shape[0] == 0:
            logging.error("未读取到任何图片数据")
            return None, 0.0
            
    except Exception as e:
        logging.error(f"数据加载失败: {e}")
        return None, 0.0
    
    # 数据分割
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
    logging.info("数据集划分为训练集和测试集。")
    
    # 定义集成方法（使用缩放和更大迭代次数以避免收敛警告）
    ensemble_methods = {
        'logistic_regression': make_pipeline(
            StandardScaler(),
            LogisticRegression(random_state=42, max_iter=5000, solver='saga', n_jobs=-1)
        ),
        'random_forest': RandomForestClassifier(n_estimators=100, random_state=42),
        'svm': make_pipeline(
            StandardScaler(),
            SVC(probability=True, random_state=42, kernel='rbf')
        ),
        'hard_voting': VotingClassifier(
            estimators=[
                ('lr', make_pipeline(StandardScaler(), LogisticRegression(random_state=42, max_iter=5000, solver='saga', n_jobs=-1))),
                ('rf', RandomForestClassifier(n_estimators=100, random_state=42)),
                ('svm', make_pipeline(StandardScaler(), SVC(probability=True, random_state=42, kernel='rbf'))),
            ],
            voting='hard'
        ),
        'soft_voting': VotingClassifier(
            estimators=[
                ('lr', make_pipeline(StandardScaler(), LogisticRegression(random_state=42, max_iter=5000, solver='saga', n_jobs=-1))),
                ('rf', RandomForestClassifier(n_estimators=100, random_state=42)),
                ('svm', make_pipeline(StandardScaler(), SVC(probability=True, random_state=42, kernel='rbf'))),
            ],
            voting='soft'
        ),
        'bagging': BaggingClassifier(
            estimator=DecisionTreeClassifier(random_state=42),
            n_estimators=50,
            random_state=42
        ),
        'pasting': BaggingClassifier(
            estimator=DecisionTreeClassifier(random_state=42),
            n_estimators=50,
            random_state=42,
            bootstrap=False
        ),
        'adaboost': AdaBoostClassifier(
            estimator=DecisionTreeClassifier(random_state=42),
            n_estimators=50,
            random_state=42
        ),
        'gradient_boosting': GradientBoostingClassifier(
            n_estimators=50,
            random_state=42
        )
    }
    
    # 存储结果
    results = []
    best_accuracy = 0
    best_model_name = ""
    best_model = None
    
    # 训练所有模型
    for name, model in ensemble_methods.items():
        logging.info(f"{name} 开始训练...")
        
        try:
            # 训练时间
            start_time = time.time()
            model.fit(X_train, y_train)
            training_time = time.time() - start_time
            
            # 预测时间
            start_time = time.time()
            y_pred = model.predict(X_test)
            prediction_time = time.time() - start_time
            
            # 计算准确率
            accuracy = accuracy_score(y_test, y_pred)
            
            # 记录结果
            results.append({
                'classifier': name,
                'training_time': training_time,
                'prediction_time': prediction_time,
                'accuracy': accuracy
            })
            
            logging.info(f"{name} 训练完成，用时 {training_time:.4f}s, 准确率 {accuracy:.4f}")
            
            # 更新最佳模型
            if accuracy > best_accuracy:
                best_accuracy = accuracy
                best_model_name = name
                best_model = model
                
        except Exception as e:
            logging.error(f"训练 {name} 时出错: {e}")
            continue
    
    # 打印结果表格 - 按照要求格式
    print("\n" + "="*85)
    print(f"{'Classifier':<20} {'Training Time (s)':<18} {'Prediction Time (s)':<20} {'Accuracy':<10}")
    print("="*85)
    
    display_order = [
        'logistic_regression', 'random_forest', 'svm', 
        'hard_voting', 'soft_voting', 'bagging', 'pasting',
        'adaboost', 'gradient_boosting'
    ]
    
    for classifier_name in display_order:
        result = next((r for r in results if r['classifier'] == classifier_name), None)
        if result:
            print(f"{result['classifier']:<20} {result['training_time']:<18.5f} {result['prediction_time']:<20.5f} {result['accuracy']:<10.5f}")
    
    print("="*85)
    print(f"最佳模型: {best_model_name}, 准确率: {best_accuracy:.5f}")
    
    # 保存最佳模型
    if best_model is not None:
        with open("best_ensemble_model.pkl", "wb") as f:
            pickle.dump(best_model, f)
        logging.info(f"最佳集成模型已保存为 best_ensemble_model.pkl，准确率: {best_accuracy:.5f}")
        
        # 保存模型信息
        model_info = {
            'model_name': best_model_name,
            'accuracy': best_accuracy,
            'feature_method': 'flat',
            'input_shape': (32, 32)
        }
        with open("model_info.pkl", "wb") as f:
            pickle.dump(model_info, f)
            
    else:
        logging.error("没有成功训练任何模型")
    
    return best_model, best_accuracy

if __name__ == '__main__':
    logging.info("开始训练集成学习模型...")
    best_model, best_accuracy = train_ensemble_models()
    
    if best_model is not None:
        logging.info("训练完成！")
    else:
        logging.error("训练失败！")